import os
from sklearn.model_selection import StratifiedKFold
from joblib import dump, load


class MyFoldData:
    def __init__(self, x, y, fold_dict_filapath="./fold_dict.gz"):
        self.x = x
        self.y = y
        self.fold_dict_filepath = fold_dict_filapath

    def create_fold_dict(self, n_splits=5):
        skf = StratifiedKFold(n_splits, shuffle=True)
        fold_dict = {}
        for index, (train_idx, val_idx) in enumerate(skf.split(self.x, self.y)):
            fold_dict[index] = (train_idx, val_idx)
        dump(fold_dict, self.fold_dict_filepath)
        return fold_dict

    def get_fold_dict(self, n_splits=5):
        if not os.path.exists(self.fold_dict_filepath):
            fold_dict = self.create_fold_dict(n_splits)
        else:
            fold_dict = load(self.fold_dict_filepath)
            if len(fold_dict) !=  n_splits:
                fold_dict = self.create_fold_dict(n_splits)
        return fold_dict

    def get_train_val_data(self, fold_idx, n_splits=5):
        fold_dict = self.get_fold_dict(n_splits)
        train_idx, val_idx = fold_dict[fold_idx]
        train_x = [self.x[idx] for idx in train_idx]
        train_y = [self.y[idx] for idx in train_idx]
        val_x = [self.x[idx] for idx in val_idx]
        val_y = [self.y[idx] for idx in val_idx]
        return train_x, train_y, val_x, val_y
